#!/usr/bin/env python
# -*- coding:utf-8 -*-
'''
@File   :   calib_rt.py
@Author :   Song
@Time   :   2023/10/12 10:34
@Contact:   songjian2022@suda.edu.cn
@intro  :   拟合预测rt和实际rt
'''
from pathlib import Path
import statsmodels.api as sm

import numpy as np
import networkx as nx
import matplotlib.pyplot as plt
from scipy.interpolate import UnivariateSpline

def generate_noise(path, noise_prop, seed):
    data = np.load(path)
    x, y = data["x"], data["y"]
    datamatrix = np.concatenate((x[:, np.newaxis], y[:, np.newaxis]), axis=1)

    datasize = len(x)
    datasize_noise = datasize * noise_prop

    noise_xrange_left, noise_xrange_right = np.min(x), np.max(x)
    noise_yrange_low, noise_yrange_up = np.min(y), np.max(y)

    np.random.seed(seed)
    noisematrix = np.concatenate(( \
        np.random.uniform(noise_xrange_left, noise_xrange_right,
                          (datasize_noise, 1), ),
        np.random.uniform(noise_yrange_low, noise_yrange_up,
                          (datasize_noise, 1), ),), axis=1)

    datamatrix = np.concatenate((datamatrix, noisematrix), axis=0)

    return datamatrix[:, 0], datamatrix[:, 1]


def o_screen_by_hist(x_data, y_data, bins):
    # 根据hist在x/y找高密度点
    extenti = (x_data.min(), x_data.max())
    extentj = (y_data.min(), y_data.max())
    hist, *edges = np.histogram2d(x_data, y_data, bins=bins,
                                  range=(extenti, extentj))
    hist_center_idx = hist.argmax(axis=1)
    cell_counts = hist[np.arange(len(hist)), hist_center_idx]
    edges_x, edges_y = edges[0], edges[1]
    edges_x = (edges_x[:-1] + edges_x[1:]) / 2
    edges_y = (edges_y[:-1] + edges_y[1:]) / 2
    x_hist = edges_x
    y_hist = edges_y[hist_center_idx]
    x_hist = x_hist[hist_center_idx > 0]  # 去掉空的cell
    y_hist = y_hist[hist_center_idx > 0]
    cell_counts = cell_counts[hist_center_idx > 0]

    return x_hist, y_hist, cell_counts

def screen_by_hist(x_data, y_data, bins):
    '''
    zhang yichi 10-14 : 实现从x和y两方向寻照hits
    '''
    # 根据hist在x/y找高密度点
    extenti = (x_data.min(), x_data.max())
    extentj = (y_data.min(), y_data.max())
    hist, edges_x, edges_y = np.histogram2d(x_data, y_data, bins=bins,
                                  range=(extenti, extentj))
    edges_x = (edges_x[:-1] + edges_x[1:]) / 2
    edges_y = (edges_y[:-1] + edges_y[1:]) / 2
    
    cell_idXy = np.stack((np.arange(hist.shape[0]),hist.argmax(axis=1)),axis=-1) # x -> max(y)
    cell_idYx = np.stack((hist.argmax(axis=0),np.arange(hist.shape[1])),axis=-1) # y -> max(x)
    cell_idxy = np.vstack((cell_idXy,cell_idYx))
    cell_idxy = np.unique(cell_idxy,axis=0) # 去重
    cell_idxy = cell_idxy[~(cell_idxy == 0).any(axis=1)] # 删去0索引项
    
    x_hist = edges_x[cell_idxy[:,0]]
    y_hist = edges_y[cell_idxy[:,1]]
    cell_counts = hist[tuple(cell_idxy.T)]

    return x_hist, y_hist, cell_counts

def o_screen_by_graph(x_screen1, y_screen1, rho):
    # 构建图，寻找最长递增路径
    G = nx.DiGraph()
    for i in range(len(x_screen1)):
        x_curr, y_curr = x_screen1[i], y_screen1[i]
        candidates_idx = (x_screen1 >= x_curr) & (y_screen1 >= y_curr)
        candidates_idx[i] = False # 防止自成环
        
        x_candidates = x_screen1[candidates_idx]
        y_candidates = y_screen1[candidates_idx]
        candidates = [((x_curr, y_curr), (x, y)) for x, y in
                      zip(x_candidates, y_candidates)]
        G.add_edges_from(candidates)

    # 从所有路径中找到最长，可能存在相等路径
    # TODO: 相等长度路径处理
    longest_path = nx.dag_longest_path(G)
    x_screen2, y_screen2 = zip(*longest_path)
    x_screen2, y_screen2 = np.array(x_screen2), np.array(y_screen2)

    return x_screen2, y_screen2

def screen_by_graph(x_screen1, y_screen1, rho):
    '''
    zhang yichi 10-14 : 增加权重 rho 为单位格中的样本点计数
    '''
    # 构建图，寻找最长递增路径
    G = nx.DiGraph()
    for i in range(len(x_screen1)):
        x_curr, y_curr = x_screen1[i], y_screen1[i]
        candidates_idx = (x_screen1 >= x_curr) & (y_screen1 >= y_curr)
        candidates_idx[i] = False # 防止自成环
        
        x_candidates = x_screen1[candidates_idx]
        y_candidates = y_screen1[candidates_idx]
        rho_candidates = rho[candidates_idx]

        candidates = [((x_curr, y_curr), (x, y),{"rho":(rho[i]+r)/2}) 
                      for x, y ,r in zip(x_candidates, y_candidates, rho_candidates)]
        G.add_edges_from(candidates)

    # 从所有路径中找到最长，可能存在相等路径
    # TODO: 相等长度路径处理
    longest_path = nx.dag_longest_path(G,weight="rho")
    x_screen2, y_screen2 = zip(*longest_path)
    x_screen2, y_screen2 = np.array(x_screen2), np.array(y_screen2)

    return x_screen2, y_screen2

def polish_ends(x_screen2, y_screen2, tol_bins):
    # 根据相邻点分别在x/y上的跨度进一步对端点过滤
    # 且左右分开比较，防止某端异常点过大导致另一端异常点无法检出
    
    center_idx = int(len(x_screen2) / 2)

    # 左侧
    x, y = x_screen2[:center_idx], y_screen2[:center_idx]
    stepx = x[1:] - x[:-1]
    good_x = (stepx / stepx[stepx > 0].min()) < tol_bins # 5%宽
    stepy = y[1:] - y[:-1]
    good_y = (stepy / stepy[stepy > 0].min()) < tol_bins
    good_xy = good_x & good_y
    print(np.where(good_xy == False)[0])
    breaks_idx = np.where(good_xy == False)[0]
    
    break_idx = 0
    if len(breaks_idx) > 0:  # 存在断点，且断点要在左边缘
        idx = np.where(breaks_idx < len(x) * 0.25)[0]
        if len(idx) > 0:
            break_idx = breaks_idx[idx][-1] + 1
    x_left, y_left = x[break_idx:], y[break_idx:]

    # 右侧
    x, y = x_screen2[center_idx:], y_screen2[center_idx:]
    stepx = x[1:] - x[:-1]
    good_x = (stepx / stepx[stepx > 0].min()) < tol_bins # 5%宽
    stepy = y[1:] - y[:-1]
    good_y = (stepy / stepy[stepy > 0].min()) < tol_bins
    good_xy = good_x & good_y
    breaks_idx = np.where(good_xy == False)[0]
    
    break_idx = len(x)
    if len(breaks_idx) > 0:  # 存在断点，且断点要在右边缘
        idx = np.where(breaks_idx > len(x) * 0.75)[0] 
        if len(idx) > 0:
            break_idx = breaks_idx[idx][0] + 1 
    x_right, y_right = x[:break_idx], y[:break_idx]
    
    x = np.concatenate([x_left, x_right])
    y = np.concatenate([y_left, y_right])
    return x, y


def fit_by_lowess(x, y):
    # 挑选合适frac
    frac=choose_frac(x,y)

    # 拟合 by lowess
    lowess = sm.nonparametric.lowess
    y_lowess = lowess(y, x, frac)
    x_fit, y_fit = zip(*y_lowess)
    x_fit, y_fit = np.array(x_fit), np.array(y_fit)
    return x_fit, y_fit



def choose_frac(x,y,delta_k_threshold=5):
    import statsmodels.api as sm
    class InputDataError(Exception):
        def __init__(self, message) -> None:
            super().__init__(message)

    for frac in np.linspace(0,1,20)[2:]:
        
        lowess = sm.nonparametric.lowess(y, x, frac)  
        x_smoothed,y_smoothed = zip(*lowess)
        
        k_list=list()

        for i in range(0,len(lowess)-1):
            delta_x=x_smoothed[i+1]-x_smoothed[i]
            delta_y=y_smoothed[i+1]-y_smoothed[i]
            k=delta_y/delta_x 
            k_list.append(k)
        k_list=np.array(k_list)
        k_delta=k_list[1:]-k_list[:-1]

        print("frac:",frac)
        print("max delta:",max(k_delta))


        if max(k_delta)<delta_k_threshold :
            return frac
    else:
        raise InputDataError("Please check if your filtered data is approiate")



def plot_all(fdir, x0, y0, x1, y1, x2, y2, x3, y3, x_fit, y_fit):
# def plot_all(fdir, x0, y0, x1, y1, x2, y2, x3, y3):
# def plot_all(fdir, x0, y0, x1, y1, x2, y2):
    data = np.load(fdir)
    x_raw, y_raw = data["x"], data["y"]
    
    plt.figure(figsize=(16,9))
    plt.plot(x0, y0, 'o', color='#808080', label='Raw+noise', markersize=0.2)
    plt.plot(x_raw, y_raw, 'ro', color='#FFE569', label='Raw', markersize=0.4)
    plt.plot(x1, y1, 'ko', label='Hist', markersize=3.1)
    plt.plot(x2, y2, 'go', label='Graph', markersize=3.1)
    plt.plot(x3, y3, 'ro', label='Polish', markersize=3.1)

    plt.plot(x_fit, y_fit, color='red')
    title = fdir.name.split('.d')[0].split('_')[-1]
    plt.title(title)
    plt.legend()
    plt.minorticks_on()
    plt.grid(True)
    plt.savefig(f"{title}.png",dpi=300)
    # plt.show()
    plt.close()

def plot_unfit(fdir, x0, y0, x1, y1, x2, y2, x3, y3):
    data = np.load(fdir)
    x_raw, y_raw = data["x"], data["y"]
    
    plt.figure(figsize=(16,9))
    plt.plot(x0, y0, 'o', color='#808080', label='Raw+noise', markersize=0.2)
    plt.plot(x_raw, y_raw, 'o', color='#FFE569', label='Raw', markersize=0.4)
    plt.plot(x1, y1, 'ko', label='Hist', markersize=3.1)
    plt.plot(x2, y2, 'go', label='Graph', markersize=3.1)
    plt.plot(x3, y3, 'ro', label='Polish', markersize=3.1)

    title = fdir.name.split('.d')[0].split('_')[-1]
    plt.title(title)
    plt.legend()
    plt.minorticks_on()
    plt.grid(True)
    plt.savefig(f"{title}.png",dpi=300)
    # plt.show()
    plt.close()

def main():
    ws = Path('./fit_five_class')

    for fidx, fdir in enumerate(ws.rglob('*.npz')):
        # if str(fdir).find('805') == -1:
        #     continue
        print(fdir)

        # raw + noise
        x0, y0 = generate_noise(fdir, noise_prop=20, seed=233)

        # screens
        x1, y1, rho= screen_by_hist(x0, y0, bins=100)
        x2, y2 = screen_by_graph(x1, y1, rho)
        x3, y3 = polish_ends(x2, y2, tol_bins=10)

        # fit and plot
        try:
            x_fit, y_fit = fit_by_lowess(x3, y3)
            plot_all(fdir, x0, y0, x1, y1, x2, y2, x3, y3, x_fit, y_fit)
        except:
            plot_unfit(fdir, x0, y0, x1, y1, x2, y2, x3, y3)

        
        # plot_all(fdir, x0, y0, x1, y1, x2, y2)


if __name__ == '__main__':
    main()
